"""
Description: RPC Server Thread for Qt app.

Author: Jin Liu
Date: 2024/05/23
"""

import sys
from typing import Callable
from qtpy import QtCore
from threading import Lock
from xmlrpc.server import SimpleXMLRPCRequestHandler, SimpleXMLRPCDispatcher
from xmlrpc.client import ServerProxy, _Method, dumps, loads, Fault
from socketserver import TCPServer
import numpy as np
from ezgl import Matrix4x4


# --- 自定义序列化和反序列化函数 ---
def serialize(params: tuple):
    """
    序列化自定义类型, 例如 Matrix4x4, np.ndarray

    将自定义类型序列化为字典, 例如 { "__Matrix4x4": Matrix4x4.toVector7d() }
    """
    if params is None:
        return params
    params = list(params)
    for i, p in enumerate(params):
        if isinstance(p, Matrix4x4):
            params[i] = {"__Matrix4x4": p.array.tolist()}
        elif isinstance(p, np.ndarray):
            params[i] = {"__ndarray": p.tolist()}

    return tuple(params)

def deserialize(params: tuple):
    """
    反序列化自定义类型, 例如 Matrix4x4, np.ndarray
    """
    if params is None:
        return params
    params = list(params)
    for i, p in enumerate(params):
        if isinstance(p, dict):
            if "__Matrix4x4" in p:
                params[i] = Matrix4x4(p["__Matrix4x4"])
            elif "__ndarray" in p:
                params[i] = np.array(p["__ndarray"])
    return tuple(params)

# ---

class QuietXMLRPCRequestHandler(SimpleXMLRPCRequestHandler):
    """
    用于抑制 XMLRPC 服务器的日志输出
    """
    def log_request(self, *args, **kwargs):
        pass  # su



class CustomRpcDispatcher(SimpleXMLRPCDispatcher):

    def _marshaled_dispatch(self, data, dispatch_method = None, path = None):
        """从序列化的数据中调度一个 XML-RPC 方法"""

        try:
            params, method = loads(data, use_builtin_types=self.use_builtin_types)
            params = deserialize(params)

            # generate response
            if dispatch_method is not None:
                response = dispatch_method(method, params)
            else:
                response = self._dispatch(method, params)
            # wrap response in a singleton tuple
            response = (response,)
            response = dumps(serialize(response), methodresponse=1,
                             allow_none=self.allow_none, encoding=self.encoding)
        except Fault as fault:
            response = dumps(fault, allow_none=self.allow_none,
                             encoding=self.encoding)
        except:
            # report exception back to server
            exc_type, exc_value, exc_tb = sys.exc_info()
            try:
                response = dumps(
                    Fault(1, "%s:%s" % (exc_type, exc_value)),
                    encoding=self.encoding, allow_none=self.allow_none,
                    )
            finally:
                exc_type = exc_value = exc_tb = None

        return response.encode(self.encoding, 'xmlcharrefreplace')


class RpcServer(TCPServer, CustomRpcDispatcher):
    """
    XML-RPC 服务端
    """

    allow_reuse_address = True

    def __init__(self, addr, requestHandler=SimpleXMLRPCRequestHandler,
                 logRequests=True, allow_none=False, encoding=None,
                 bind_and_activate=True, use_builtin_types=False):
        self.logRequests = logRequests

        CustomRpcDispatcher.__init__(self, allow_none, encoding, use_builtin_types)
        TCPServer.__init__(self, addr, requestHandler, bind_and_activate)


class RpcClient(ServerProxy):
    """
    XML-RPC 客户端
    """
    def __init__(self, uri, transport=None, encoding=None, verbose=False,
                 allow_none=False, use_datetime=False, use_builtin_types=False,
                 *, headers=(), context=None):
        super().__init__(uri, transport, encoding, verbose, allow_none, use_datetime, use_builtin_types, headers=headers, context=context)
        self.mutex = Lock()

    def request(self, methodname, params):
        with self.mutex:  # 保证线程安全
            # call a method on the remote server
            request = dumps(serialize(params), methodname, encoding=self._ServerProxy__encoding,
                            allow_none=self._ServerProxy__allow_none).encode(self._ServerProxy__encoding, 'xmlcharrefreplace')

            response = self._ServerProxy__transport.request(
                self._ServerProxy__host, self._ServerProxy__handler, request,
                verbose=self._ServerProxy__verbose
            )

            response = deserialize(response)
            if len(response) == 1:
                response = response[0]

        return response

    def __getattr__(self, name):
        # magic method dispatcher
        return _Method(self.request, name)


class RpcServerThread(QtCore.QThread):
    """
    RPC Server Thread for Qt app.
    """
    def __init__(self, parent=None, port=24123, quiet=True):
        """
        Initialization.

        Parameters:
        - parent : QWidget, optional, default: None
        - port : int, optional, default: 24123, 端口号
        - quiet : bool, optional, default: True, 是否抑制日志输出
        """
        super().__init__(parent)
        self.port = port
        Handler = QuietXMLRPCRequestHandler if quiet else SimpleXMLRPCRequestHandler
        self.server = RpcServer(("localhost", port),
                                    requestHandler=Handler,
                                    allow_none=True)
        self.server.register_function(print, "_echo")  # for test

    def register_function(self, func:Callable, name:str):
        """
        注册函数
        """
        self.server.register_function(func, name)

    def run(self):
        self.server.serve_forever()
